#include "formation_avoidance/distribute_optimizer.h"

bool DistributeOptimizer::getGradient()
{
	Eigen::Vector4f Gradient=Eigen::Vector4f::Zero();

	current_v=costJ(variable,true,1000);
	double delta=0.00001;

	if(OptimizerID<=1)
	{
		for(int i=OptimizerID;i<=OptimizerID+2;i++)
		{
			Eigen::Vector4f tmp=variable;
			tmp(i)+=delta;
			double tmp_v=0;
			tmp_v=costJ(tmp,false,1000);
			Gradient(i)=(tmp_v-current_v)/delta;
		}
	}
	else if(OptimizerID==2)
	{
		int index[2]={2,3};
		for(int i=0;i<=1;i++)
		{
			Eigen::Vector4f tmp=variable;
			tmp(index[i])+=delta;
			double tmp_v=costJ(tmp,false,1000);
			Gradient(index[i])=(tmp_v-current_v)/delta;
		}
	}
	else if(OptimizerID==3)
	{
		int index[2]={0,3};
		for(int i=0;i<=1;i++)
		{
			Eigen::Vector4f tmp=variable;
			tmp(index[i])+=delta;
			double tmp_v=costJ(tmp,false,1000);
			Gradient(index[i])=(tmp_v-current_v)/delta;
		}
	}

	gradient=Gradient;

	NeighbourVariable[OptimizerID]=variable;
	NeighbourGradient[OptimizerID]=gradient;


	return true;
}

bool DistributeOptimizer::sendVariableAndReceiving()
{
	othergetflag=0;
	megetflag=0;

	for(int i=0;i<4;i++)
	{
		selfgetflag[i]=0;
		neighbourgetflag[i]=0;
	}
	selfgetflag[OptimizerID]=1;
	neighbourgetflag[OptimizerID]=1;

	formation_avoidance::DistributedOptimize DistributedOptimizeMsg;

	for(int i=0;i<4;i++)
    {
		DistributedOptimizeMsg.variable[i]=variable(i);
		DistributedOptimizeMsg.gradient[i]=gradient(i);
		DistributedOptimizeMsg.y[i]=y(i);
		DistributedOptimizeMsg.selfgetflag[i]=selfgetflag[i];
    }
	DistributedOptimizeMsg.OptimizeID=OptimizerID;
	DistributedOptimizeMsg.send_count=optimize_round;
	DistributedOptimizeMsg.megetflag=megetflag;

    // trajectory_msgs::JointTrajectoryPoint optimize_msg;
	// for(int i=0;i<4;i++)
    // {
	// 	// ROS_INFO("variable(I): %f",variable(i));
    //     optimize_msg.positions.push_back(variable(i));  
    //     optimize_msg.positions.push_back(gradient(i));  
	// 	optimize_msg.positions.push_back(y(i));  

    // }
    // optimize_msg.velocities.push_back(OptimizerID);
	// optimize_msg.velocities.push_back(optimize_round);

	// for(int i=0;i<4;i++)
	// {
	// 	optimize_msg.accelerations.push_back(selfgetflag[i]);
	// }
	// optimize_msg.effort.push_back(megetflag);
	// optimize_msg.effort.push_back(0);

	int count=0;
	int send_count=0;
	while(1)
	{	
		ros::spinOnce();
		send_count++;

		for(int i=0;i<4;i++)
		{
			DistributedOptimizeMsg.selfgetflag[i]=selfgetflag[i];
		}
		DistributedOptimizeMsg.megetflag=megetflag;

		DistributedOptimizeMsg.Header.stamp=ros::Time::now();
		variable_gradient_pub.publish(DistributedOptimizeMsg);

		// for(int i=0;i<4;i++)
		// {
		// 	optimize_msg.accelerations[i]=selfgetflag[i];
		// }

		// optimize_msg.effort[0]=megetflag;
		// optimize_msg.effort[1]=send_count;

		// ROS_INFO("id %d optimize_msg 0: %f 1: %f optimize_msg.effort %f",OptimizerID,optimize_msg.accelerations[0],optimize_msg.accelerations[1],optimize_msg.effort[0]);

		// ROS_INFO("ID : %d megetflag: %d",OptimizerID,megetflag);

		// ROS_INFO_THROTTLE(1,"neighbourgetflag: %d %d %d %d",neighbourgetflag[0],neighbourgetflag[1],neighbourgetflag[2],neighbourgetflag[3]);
		if(neighbourgetflag[0]+neighbourgetflag[1]+neighbourgetflag[2]+neighbourgetflag[3]<4)
		// if(neighbourgetflag[0]+neighbourgetflag[1]+neighbourgetflag[2]<3)
		// if(neighbourgetflag[0]+neighbourgetflag[1]<2)
		// if(othergetflag==0)
		{
			// optimize_msg.time_from_start=ros::Duration(ros::Time::now().toSec());
			// // ROS_INFO("OPTIMIZE id %d send message",OptimizerID);
			// ROS_INFO_THROTTLE(0.5,"\033[1;%dm id %d send message send_count: %d \033[0m",OptimizerColorID,OptimizerID,send_count);
			// variable_gradient_pub.publish(optimize_msg);

			// DistributedOptimizeMsg.Header.stamp=ros::Time::now();
			// variable_gradient_pub.publish(DistributedOptimizeMsg);

		}
		else{

			// ROS_WARN("optimizeid %d neighbourgetflag: %d %d %d %d",OptimizerID,neighbourgetflag[0],neighbourgetflag[1],neighbourgetflag[2],neighbourgetflag[3]);
			while(count<=10)
			{	
				// variable_gradient_pub.publish(optimize_msg);
				DistributedOptimizeMsg.Header.stamp=ros::Time::now();
				variable_gradient_pub.publish(DistributedOptimizeMsg);
				// loop_rate.sleep();
				ros::Duration(0.001).sleep();
				count++;
			}
			// ROS_INFO("OPTIMIZEID %d SEND COUNT: %d",OptimizerID,send_count);
			return true;
		}
		// ROS_INFO_THROTTLE(0.5,"\033[1;%dm id:%d ros::spinonce  \033[0m",OptimizerColorID,OptimizerID);
		ros::spinOnce();
		loop_rate.sleep();
	}
}


bool DistributeOptimizer::updateVariable()
{
	if(!GDInitFlag) return true;

	Eigen::Vector4f update_variable=Eigen::Vector4f::Zero();
	for(int i=0;i<4;i++)
	{
		update_variable+=W(OptimizerID,i)*NeighbourVariable[i];
	}

	//backtraking line search
	double rou=0.8;
	double c=0.3;
	double gama=0.8;
	Eigen::Vector4f p=-y;
	while(costJ(variable+gama*p,true,1000)>current_v+c*gama*y.transpose()*p)
	{
		gama=rou*gama;
		// if(gama<1e-6)break;
	}
	// ROS_INFO("\033[1;%dm id:%d,gama:%f\033[0m",OptimizerColorID,OptimizerID,gama);
	update_variable-=gama*y;

	variable=update_variable;
	variable_mapping();

	return true;
}

void DistributeOptimizer::updatey()
{
	if(!GDInitFlag)
	{
		y=gradient;
		Neighbour_y[OptimizerID]=y;
		lastgradient=gradient;

		GDInitFlag=true;
		return ;
	}	

	Eigen::Vector4f update_y=Eigen::Vector4f::Zero();
	for(int i=0;i<4;i++)
	{
		update_y+=W(OptimizerID,i)*Neighbour_y[i];
	}
	update_y+=gradient-lastgradient;
	y=update_y;
	Neighbour_y[OptimizerID]=y;

	lastgradient=gradient;

}
void DistributeOptimizer::gradientDescent()
{
	ros::Time start=ros::Time::now();
	double communicationtime=0;
	double updatetime=0;
	ros::Time startime=ros::Time::now();
	ros::Time endtime=ros::Time::now();

	for(int i=0;i<200;i++)
	{
		startime=ros::Time::now();
		updateVariable();
		endtime=ros::Time::now();
		ROS_INFO("\033[1;%dm id %d Updated Variable %f %f %f %f current_v :%f updatevariable spend:%f round:%d\033[0m",OptimizerColorID,OptimizerID,variable(0),variable(1),variable(2),variable(3),current_v,(endtime-startime).toSec(),optimize_round);
		updatetime+=(endtime-startime).toSec();

		startime=ros::Time::now();
		getGradient();
		endtime=ros::Time::now();

		// ROS_INFO("\033[1;%dm ID: %d calculate GRADIENT %f %f %f %f current_v:%f  gradientspend:%f \033[0m",OptimizerColorID,OptimizerID,gradient(0),gradient(1),gradient(2),gradient(3),current_v,(endtime-startime).toSec());
		
		startime=ros::Time::now();
		updatey();
		endtime=ros::Time::now();
		// ROS_INFO("\033[1;%dm ID: %d updateyspend:%f \033[0m",OptimizerColorID,OptimizerID,(endtime-startime).toSec());

		startime=ros::Time::now();
		sendVariableAndReceiving();
		endtime=ros::Time::now();
		// ROS_INFO("\033[1;%dm optimizerid: %d time spend communication: %f optimize_round: %d \033[0m",OptimizerColorID,OptimizerID,(endtime-startime).toSec(),optimize_round);
		communicationtime+=(endtime-startime).toSec();
		optimize_round++;

		bool GDendflag=true;
		for(int j=0;j<3;j++)
		{
			if((NeighbourVariable[j]-NeighbourVariable[j+1]).norm()>0.0001||(Neighbour_y[0]+Neighbour_y[1]+Neighbour_y[2]+Neighbour_y[3]).norm()>0.0001)
			{
				GDendflag=false;
				break;
			}
		}
		if(GDendflag)break;

		// ros::Duration(1.0).sleep();
	}
	ros::Time end=ros::Time::now();

	ROS_INFO("\033[1;%dm ID %d finish GD spend %f optimize_round:%d variable:%f %f %f %f communication time:%f updatetime:%f\033[0m  ",OptimizerColorID,OptimizerID,(end-start).toSec(),optimize_round
	,variable(0),variable(1),variable(2),variable(3),communicationtime,updatetime);

}


void DistributeOptimizer::optimize_sever_cb(const formation_avoidance::DistributedOptimize::ConstPtr& msg)
{
	// if(msg->send_count!=optimize_round) return;
	// int getID=msg->OptimizeID;
	// // ROS_INFO("\033[1;%dm id:%d getid:%d \033[0m",OptimizerColorID,OptimizerID,getID);

	// // ROS_INFO("msg->accelerations: %f %f %f %f",msg->accelerations[0],msg->accelerations[1],msg->accelerations[2],msg->accelerations[3]);
	// if(getID==OptimizerID) 
	// {
	// 	ROS_WARN("RETURN-------------");
	// 	return;
	// }
	// megetflag=1;
	// selfgetflag[getID]=1;
	// for(int i=0;i<4;i++)
	// {
	// 	NeighbourVariable[getID](i)=msg->variable[i];
	// 	NeighbourGradient[getID](i)=msg->gradient[i];
	// 	Neighbour_y[getID](i)=msg->y[i];

	// }

	// neighbourgetflag[getID]=msg->selfgetflag[OptimizerID];

	// othergetflag=msg->megetflag;
	// int send_count=msg->send_count;


	// // ROS_INFO("id %d neighbourgetflag[getid]: %d  othergetflag: %d",OptimizerID,neighbourgetflag[getID],othergetflag);

	// // ROS_INFO("ID: %d get message, othergetflag : %d send_count: %d",OptimizerID,othergetflag,send_count);

	if(fabs(msg->round-round)>=10)
	{
		// ROS_ERROR("FATAL ERROR");
		return;
	}

	if(msg->best<globalBestParticle.second)
	{
		globalBestParticle.second=msg->best;
		globalBestParticle.first(0)=msg->particle[0];
		globalBestParticle.first(1)=msg->particle[1];
		globalBestParticle.first(2)=msg->particle[2];
		globalBestParticle.first(3)=msg->particle[3];
	}

}


void particle_mapping(Eigen::Vector4f& variable)
{
	double PI=3.1415926;
	for(int i=0;i<4;i++)
	{
		while(variable(i)>=PI)
			variable(i)-=2*PI;
		while(variable(i)<-PI)
			variable(i)+=2*PI;
	}

}


void DistributeOptimizer::PSO()
{
	ros::Time start=ros::Time::now();
	//init
	srand((ros::Time::now().toSec())/(OptimizerID+0.5));

    for(int i=0;i<PARTICLE_NUM;i++)
    {
		for(int k=0;k<4;k++)
		{
			double max=3.1;
			double min=-3.1;
			particle[i](k)=(max-min)*rand()/(RAND_MAX)+min;
		}
		particle_mapping(particle[i]);
		localBestParticle[i].second=9999;

		for(int k=0;k<4;k++)
		{
			double max=0.2;
			double min=-0.2;
			particle_v[i](k)=(max-min)*rand()/(RAND_MAX)+min;
		}

    }
    globalBestParticle.second=9999;
	// ROS_INFO("id:%d, init particle[0]:%f %f %f %f",OptimizerID,particle[0](0),particle[0](1),particle[0](2),particle[0](3));
	//init finish


    double c1=2;
    double c2=2;
    int t_max=40;
    Eigen::Vector4f variable10;
    for(int t=0;t<t_max;t++)
    {

		for(int i=0;i<PARTICLE_NUM;i++)
		{
			//calculate min

			double tmpmin=costJ(particle[i],true,9999);
			
			if(tmpmin<localBestParticle[i].second)
			{
				// ROS_INFO("id: %d AAAAAAAAAA",OptimizerID);
				localBestParticle[i].first=particle[i];
				localBestParticle[i].second=tmpmin;
			}

			if(tmpmin<globalBestParticle.second)
			{
				// ROS_INFO("id: %d jjjjjjjjjjjjjjjjj",OptimizerID);
				globalBestParticle.first=particle[i];
				globalBestParticle.second=tmpmin;
			}
		}



		for(int i=0;i<PARTICLE_NUM;i++)
		{
			//calculate v
			double w_min=0.2;
			double w_max=0.9;

			double w=w_max-(w_max-w_min)/t_max*t;
			particle_v[i]=w*particle_v[i]+c1*rand()/(RAND_MAX)*(localBestParticle[i].first-particle[i])+c2*rand()/(RAND_MAX)*(globalBestParticle.first-particle[i]);

			double vmax=0.3;
			for(int j=0;j<4;j++)
			{
				if(fabs(particle_v[i][j])>vmax)
				{
					particle_v[i][j]=particle_v[i][j]/fabs(particle_v[i][j])*vmax;
				}
			}
			//update particle

			particle[i]+=particle_v[i];
			particle_mapping(particle[i]);



        }

		formation_avoidance::DistributedOptimize optimize_msg;
		optimize_msg.Header.stamp=ros::Time::now();
		optimize_msg.best=globalBestParticle.second;
		for(int i=0;i<4;i++)
			optimize_msg.particle[i]=globalBestParticle.first(i);
		optimize_msg.round=round;
		variable_gradient_pub.publish(optimize_msg);

		ros::spinOnce();

        if(t==10)variable10=globalBestParticle.first;

        
        // ROS_INFO("globalbestVARIABLE: %f %f %f %f value:%f count:%d",globalBestParticle.first(0),globalBestParticle.first(1),globalBestParticle.first(2),globalBestParticle.first(3),globalBestParticle.second
        // ,t+1);

    }
	variable=globalBestParticle.first;
	ros::Time end=ros::Time::now();


	// ROS_INFO("\033[1;%dm ID %d VARIABLE: %f %f %f %f value:%f time:%f \033[0m",OptimizerColorID,OptimizerID,globalBestParticle.first(0),globalBestParticle.first(1),globalBestParticle.first(2),globalBestParticle.first(3),globalBestParticle.second
	// ,(end-start).toSec());
}

void DistributeOptimizer::initial_variable(Eigen::Vector4f init_variable)
{
	variable=init_variable;
	variable_mapping();
	y=Eigen::Vector4f::Zero();
	GDInitFlag=false;
	optimize_round=0;

	double average=1/4.0;
    W<<average,average,average,average,
       average,average,average,average,
       average,average,average,average,
       average,average,average,average;

	for(int i=0;i<4;i++)
	{
		selfgetflag[i]=0;
		neighbourgetflag[i]=0;
		NeighbourVariable[i]=Eigen::Vector4f::Zero();
		NeighbourGradient[i]=Eigen::Vector4f::Zero();
	}

	gradient=Eigen::Vector4f::Zero();

	selfgetflag[OptimizerID]=1;
	neighbourgetflag[OptimizerID]=1;


}

void DistributeOptimizer::variable_mapping()
{
	double PI=3.1415926;
	for(int i=0;i<4;i++)
	{
		while(variable(i)>=PI)
			variable(i)-=2*PI;
		while(variable(i)<-PI)
			variable(i)+=2*PI;
	}
}


void DistributeOptimizer::init(ros::NodeHandle &nh)
{
	ROS_INFO("DistributeOptimizer INIT");

    std::ostringstream str_iris;
    str_iris<<"/agent"<<OptimizerID;
	// variable_gradient_pub=nh.advertise<trajectory_msgs::JointTrajectoryPoint>(str_iris.str()+"/optimize_server", 1, true);
	variable_gradient_pub=nh.advertise<formation_avoidance::DistributedOptimize>(str_iris.str()+"/optimize_server", 1, true);

	int subcount=0;
	for(int i=0;i<4;i++)
	{
		if(i==OptimizerID)continue;
		std::ostringstream str_iris;
		str_iris<<"/agent"<<i;
		variable_gradient_sub[subcount]=nh.subscribe(str_iris.str()+"/optimize_server", 6, &DistributeOptimizer::optimize_sever_cb,this);
		subcount++;
	}
	for(int i=0;i<4;i++)
	{
		selfgetflag[i]=0;
		neighbourgetflag[i]=0;
		NeighbourVariable[i]=Eigen::Vector4f::Zero();
		NeighbourGradient[i]=Eigen::Vector4f::Zero();
	}

	variable=Eigen::Vector4f::Zero();
	gradient=Eigen::Vector4f::Zero();

	selfgetflag[OptimizerID]=1;
	neighbourgetflag[OptimizerID]=1;


}